from pypatnlp import *
from random import sample
import os

SAMPLE_SIZE = 1

TASK_PATH = os.path.join('tasks', 'autorituvastus')
DATA_PATH = os.path.join(TASK_PATH, 'data')
REFERENCE_PATH = os.path.join(TASK_PATH, 'reference')
TRAINTEST_PATH = os.path.join(TASK_PATH, 'traintest')

def get_paths():
    return [name for name in os.listdir(DATA_PATH) if name.endswith('.pycorp')]

# get the references for all pycorps
pycorps = []
names = []
for path in get_paths():
    print path
    pycorps.append(PyCorpus(os.path.join(DATA_PATH, path), readonly=True))
    names.append(os.path.basename(path))

for idx in range(len(names)):
    print 'Processing ', names[idx]
    name = names[idx]
    # make training corpus
    doc_ids = list(pycorps[idx].keys())
    traincorp = PyCorpus(os.path.join(TRAINTEST_PATH, 'train_' + name))
    traincorp.autocommit(False)
    for doc_id in doc_ids[:-1]:
        traincorp[doc_id] = pycorps[idx][doc_id]
    traincorp.commit()
    traincorp.close()
    
    test_candidates = []
    for idx2 in range(len(names)):
        if idx == idx2:
            continue
        for doc_id in pycorps[idx2].keys():
            test_candidates.append((idx2, doc_id))
    test_candidates = sample(test_candidates, SAMPLE_SIZE)
    
    testcorp = PyCorpus(os.path.join(TRAINTEST_PATH, 'test_' + name))
    testcorp.autocommit(False)
    testcorp['YES'] = pycorps[idx][doc_ids[-1]]
    for idx3, (idx2, doc_id) in enumerate(test_candidates):
        testcorp['NO' + str(idx3)] = pycorps[idx2][doc_id]
    testcorp.commit()
